
import struct
import numpy as np
from PIL import Image
from vtkmodules.vtkCommonCore import VTK_SHORT, VTK_UNSIGNED_CHAR
from vtkmodules.vtkCommonDataModel import vtkImageData
from vtkmodules.vtkFiltersGeometry import vtkImageDataGeometryFilter
from vtkmodules.vtkCommonCore import VTK_SHORT
from vtkmodules.vtkRenderingCore import (
    vtkActor,
    vtkPolyDataMapper,
    vtkTexture,
)
from vtkmodules.vtkInteractionStyle import *
from vtkmodules.vtkFiltersGeneral import vtkWarpScalar
from vtkmodules.vtkRenderingLOD import vtkLODActor
from vtkmodules.util import numpy_support
from visualizationCom.tools.color import hex_to_rgba

pixel_spacing = [0.09, 0.006]  # x, y
depth_scaling = 0.0006


def imread_16bitDepthBMP(path):
    '''
    struct BMPHeader {
        uint16_t signature;   // 文件标识 BM
        uint32_t fileSize;    // 文件大小
        uint32_t reserved;    // 保留字段
        uint32_t dataOffset;  // 数据偏移量
        uint32_t headerSize;  // 文件头大小
        uint32_t width;       // 图像宽度
        uint32_t height;      // 图像高度
        uint16_t planes;      // 颜色平面数
        uint16_t bitCount;    // 每像素位数
        uint32_t compression; // 压缩方式
        uint32_t imageSize;   // 图像大小
        uint32_t xPixelsPerMeter; // 水平分辨率
        uint32_t yPixelsPerMeter; // 垂直分辨率
        uint32_t colorsUsed;      // 使用的颜色数
        uint32_t colorsImportant; // 重要颜色数
    };
    '''

    file = open(path, "rb")
    header = file.read(54)
    width, height = struct.unpack_from('II', header, 18)
    # read binary data (short)
    depth_data = file.read(width * height * 2)
    # convert to vtk array
    vtk_array = numpy_support.numpy_to_vtk(
        np.frombuffer(depth_data, dtype=np.short),
        deep=True, array_type=VTK_SHORT)

    imageData = vtkImageData()
    imageData.SetDimensions(width, height, 1)
    imageData.SetSpacing(pixel_spacing[0], pixel_spacing[1], 0)
    imageData.SetOrigin(0, 0, 0)
    imageData.GetPointData().SetScalars(vtk_array)
    return imageData


def imread_Label(path):
    # read with PIL
    im = Image.open(path)
    # im = im.convert("RGB")
    colors = im.getcolors()
    print(colors, '原图缺陷颜色')

    # colors_temp = [color[1] for color in colors]

    # for index, value in enumerate(colors_temp):
    #     rgb_list = list(value)
    #     print(rgb_list,'单一颜色')
    #     colorsArry.append(rgb_list)

    labels = np.flip(np.array(im), axis=0)
    # vtk_array = numpy_support.numpy_to_vtk(indices.ravel(), deep=True)
    return labels, len(colors)


class LabelColorMapper():
    def __init__(self) -> None:
        self.image = vtkImageData()
        self.image.SetDimensions(1, 1, 1)
        self.image.AllocateScalars(VTK_UNSIGNED_CHAR, 4)

        self.texture = vtkTexture()
        self.texture.SetInputDataObject(self.image)
        # self.texture.SetUseSRGBColorSpace(1)
        self.texture.InterpolateOff()
        self.texture.RepeatOff()

    def GetTexture(self):
        return self.texture

    def SetNumColors(self, n):
        self.image.SetDimensions(n, 1, 1)
        self.image.AllocateScalars(VTK_UNSIGNED_CHAR, 4)

    def SetColor(self, idx, color):
        self.image.SetScalarComponentFromFloat(idx, 0, 0, 0, color[0])
        self.image.SetScalarComponentFromFloat(idx, 0, 0, 1, color[1])
        self.image.SetScalarComponentFromFloat(idx, 0, 0, 2, color[2])
        self.image.SetScalarComponentFromFloat(idx, 0, 0, 3, 255)
        self.image.Modified()

    def FillColor(self, color):
        for idx in range(self.image.GetDimensions()[0]):
            self.SetColor(idx, color)


def GetTexCoords(label, size):
    texcoord = label.ravel().astype(np.float32) / float(size - 1)
    texcoord = np.stack([texcoord, texcoord], axis=-1)
    vtk_array = numpy_support.numpy_to_vtk(texcoord, deep=True)
    # vtk_array.SetName('TCoords')
    return vtk_array


def loadMeshOrPoint(depth, label, size, showType):
    # create surface from depth
    surface = vtkImageDataGeometryFilter()
    surface.SetInputData(depth)

    # scale z-depth
    warp = vtkWarpScalar()
    warp.SetInputConnection(surface.GetOutputPort())
    warp.SetScaleFactor(depth_scaling)
    warp.UseNormalOn()
    warp.SetNormal(0, 0, 1)
    warp.Update()

    texCoord = GetTexCoords(label, size)

    # get poly data
    warp.Update()
    polyData = warp.GetOutput()
    polyData.GetPointData().SetTCoords(texCoord)

    colorMapper = LabelColorMapper()
    colorMapper.SetNumColors(size)
    colorMapper.SetColor(0, [255, 255, 255])
    colorMapper.SetColor(1, [255, 0, 0]) # 这条数据直接写死默认选中，真实代码需要根据传过来的id进行选中
    colorMapper.SetColor(2, [0,   255,   0])
    colorMapper.SetColor(3, [0,   0,   255])
    colorMapper.SetColor(4, [128, 0, 128])
    colorMapper.SetColor(5, [255, 255, 0])
    colorMapper.SetColor(6, [0, 255, 255])
    colorMapper.SetColor(7, [255, 0, 255])

    global actor
    if showType == "Mesh":
        surfaceMapper = vtkPolyDataMapper()
        surfaceMapper.SetInputData(polyData)
        surfaceMapper.ScalarVisibilityOff()
        surfaceActor = vtkActor()
        surfaceActor.SetMapper(surfaceMapper)
        surfaceActor.SetTexture(colorMapper.GetTexture())
        actor = surfaceActor
    else:
        pointMapper = vtkPolyDataMapper()
        pointMapper.SetInputData(polyData)
        pointMapper.ScalarVisibilityOff()
        pointActor = vtkActor()
        # pointActor = vtkLODActor() 这个actor可以在交互时少渲染，停下后渲染全图，但有时候会在交互时整张图都不渲染了，所以先用上述的代码实现
        # pointActor.SetNumberOfCloudPoints(int(depth.GetNumberOfPoints()/10))
        pointActor.SetMapper(pointMapper)
        pointActor.SetTexture(colorMapper.GetTexture())
        pointActor.GetProperty().SetRepresentationToPoints()
        actor = pointActor


def load(depth_file, label_file, showType):
    # depth image
    depth = imread_16bitDepthBMP(depth_file)
    # defects mask
    global size
    label, size = imread_Label(label_file)

    loadMeshOrPoint(depth, label, size, showType)

    return actor


def selectDefect(id):
    colorMapper = LabelColorMapper()
    colorMapper.SetNumColors(size)
    colorMapper.SetColor(0, [255, 255, 255])
    colorMapper.SetColor(1, [255, 192, 203])
    colorMapper.SetColor(2, [0,   255,   0])
    colorMapper.SetColor(3, [0,   0,   255])
    colorMapper.SetColor(4, [128, 0, 128])
    colorMapper.SetColor(5, [255, 255, 0])
    colorMapper.SetColor(6, [0, 255, 255])
    colorMapper.SetColor(7, [255, 0, 255])
    actor.SetTexture(colorMapper.GetTexture())

    colorMapper.SetColor(id, [255, 0, 0])
